package com.sunzm.spark.sql.udaf

import org.apache.spark.sql.SparkSession

/**
 *
 * SparkSQL UDAF 示例程序
 *
 * @author Administrator
 * @version 1.0
 * @date 2022-01-12 20:54
 */
object SparkSQLUDAFDemo {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession
      .builder()
      .appName(this.getClass.getSimpleName.stripSuffix("$"))
      .master("local[*]")
      .config("spark.default.parallelism", 8)
      .config("spark.sql.shuffle.partitions", 8)
      .getOrCreate()

    spark.sparkContext.setLogLevel("WARN")

    import spark.implicits._

    import org.apache.spark.sql.functions._

    /*
     * 对于下面的数据
     *
     * 如果先按照区进行聚合,得到：
     * +--------+------+-----------+---+
       |province|area  |order_count|uv |
       +--------+------+-----------+---+
       |北京市  |东城区|40         |1  |
       |北京市  |海淀区|60         |2  |
       +--------+------+-----------+---+
       *
       * 然后再按照省份进行聚合,上面聚合结果中的 uv 就不能使用了, 因为uv是需要去重的, 只能根据最开始的明细数据再去聚合
       * +--------+-----------+---+
         |province|order_count|uv |
         +--------+-----------+---+
         |北京市  |100        |2  |
         +--------+-----------+---+
         *
         * 一个办法是, 在中间聚合结果中把去重的用户Id保存下来, 比如: 北京市  |海淀区|60         |2  | 12
         *
         *  但是这种方法,如果用户Id比较长,存储大量的用户Id,会占用很大的空间
         *
         * 所以, 可以考虑使用 bitmap, 基本思想是 比如总共有3个用户, 那么就有一个长度为3的数组, 某个用户出现了, 对应位置就是1, 否则就是0.
         * 需要做一个用户Id与数组中对应位置的映射,然后对应位置放0和1就行了. 实际上, 如果数组长度为 10 万, 只有个别用户出现了, 不需要真正
         * 每个位置都存上具体的数字,数组就会很稀疏. 比如可以存 (1,9000)表示第1和9000个位置是1, 再比如 11, 12, 13, 14, 15, 27, 28, 29
         * 可以压缩为两个二元组11, 4, 27, 2，表示11后面紧跟着4个连续递增的值，27后面跟着2个连续递增的值。
         * 具体实现选择了 RoaringBitmap, 它内部会随着数据的增多会自动扩容,而且使用了压缩算法,存储的数据量进一步减少
       *
     */
    val dataSeq = Seq[(Int, String, String, Int)](
      (1, "北京市", "海淀区", 10),
      (1, "北京市", "海淀区", 20),
      (2, "北京市", "海淀区", 30),
      (1, "北京市", "东城区", 20),
      (1, "北京市", "东城区", 20)
    )

    spark.createDataset(dataSeq)
      .toDF("user_id", "province", "area", "order_count")
      .createOrReplaceTempView("v_data")

    // 注册自定义UDAF
    spark.udf.register("int_to_bitmap",udaf(new IntToBitMapUDAF))
    spark.udf.register("merge_bitmap",udaf(new BitMapMergeUDAF))

    //自定义udf, 计算基数 (count(distinct 字段名))
    val card = (bmBytes:Array[Byte])=>{
      BitMapUtil.deSerBitMap(bmBytes).getCardinality
    }
    spark.udf.register("card_bm",card)

    //按照省份和区域分组,求每个省份每个区域的订单数和用户数
    // 并把去重后的userId以bitmap的形式存储到一个临时视图中,以便后面按照省份聚合的时候,可以复用
    val areaDF = spark.sql(
      """
        | SELECT
        |   province,
        |   area,
        |   sum(order_count) AS order_count,
        |   count(distinct user_id) AS uv_cnt_dist,
        |   card_bm(int_to_bitmap(user_id)) AS uv,
        |   int_to_bitmap(user_id) AS bitmap
        | FROM v_data
        | GROUP BY province, area
        |""".stripMargin)

    areaDF.createOrReplaceTempView("v_area")

    areaDF.show(100, false)


    spark.sql(
      """
        | SELECT
        |   province,
        |   sum(order_count) AS order_count,
        |   card_bm(merge_bitmap(bitmap)) AS uv
        | FROM v_area
        | GROUP BY province
        |""".stripMargin)
      .show(100, false)

    spark.stop()
  }
}
